from cProfile import label
import numpy as np
import os
from apk_process import process_apk
import pickle
import scipy.io as sio
import scipy.sparse as sparse
import time


# apk_path_list = ["./sample_apk/000db93ae6b67e6303e5b5b4839df016.apk", 
#             "./sample_apk/000f3a904966dcffed17a965c183ba70.apk",
#             "./sample_apk/000316cad2ed6477e73d839f60115e5f.apk"]
apk_path_list = []
apk_test_or_train = []
apk_labels = []

for (dirpath, dirnames, filenames) in os.walk("./sample/test"):
    for filename in filenames:
        apkname = os.path.join(dirpath, filename)
        apk_istrain = 0
        if apkname.split('/')[-2] == "adware":
            apk_label = [1, 0, 0]
        elif apkname.split('/')[-2] == "normal":
            apk_label = [0, 1, 0]
        elif apkname.split('/')[-2] == "trojan":
            apk_label = [0, 0, 1]
        apk_path_list.append(apkname)
        apk_test_or_train.append(apk_istrain)
        apk_labels.append(apk_label)

for (dirpath, dirnames, filenames) in os.walk("./sample/train"):
    for filename in filenames:
        apkname = os.path.join(dirpath, filename)
        apk_istrain = 1
        if apkname.split('/')[-2] == "adware":
            apk_label = [1, 0, 0]
        elif apkname.split('/')[-2] == "normal":
            apk_label = [0, 1, 0]
        elif apkname.split('/')[-2] == "trojan":
            apk_label = [0, 0, 1]
        apk_path_list.append(apkname)
        apk_test_or_train.append(apk_istrain)
        apk_labels.append(apk_label)

# apk_path_list = apk_path_list[0:15]
# apk_test_or_train = apk_test_or_train[0:15]
# apk_labels = apk_labels[0:15]


per_dict = {}  # permission
cla_dict = {}  # class
lib_dict = {}  # library
int_dict = {}  # interface
met_dict = {}  # method

'''
class Myapk():
    def __init__(self, permission, apk_class, library, interface, methods):
        self.permission = permission
        self.apk_class = apk_class
        self.library = library
        self.interface = interface
        self.methods = methods
        self.per_vector = []
        self.cla_vector = []
        self.lib_vector = []
        self.int_vector = []
        self.met_vector = [] 
'''

# per_feature = []    # permission
# cla_feature = []    # class
# lib_feature = []    # library
# int_feature = []    # interface
# met_feature = []    # method

# Collect all features
apk_list = []
apk_test_or_train_new = []
apk_labels_new = []
apk_path_list_new = []
count = 0
time_point_0 = time.time()
for apk_idx in range(len(apk_path_list)):
    try:
        apk = process_apk(apk_path_list[apk_idx], per_dict, cla_dict, lib_dict, int_dict, met_dict)   
        # per_feature.extend(apk.permission)
        # per_feature = list(set(per_feature))
        
        # cla_feature.extend(apk.apk_class)
        # cla_feature = list(set(cla_feature))
        
        # lib_feature.extend(apk.library)
        # lib_feature = list(set(lib_feature))
        
        # int_feature.extend(apk.interface)
        # int_feature = list(set(int_feature))
        
        # met_feature.extend(apk.methods)
        # met_feature = list(set(met_feature))
        
        apk_list.append(apk)
        apk_test_or_train_new.append(apk_test_or_train[apk_idx])
        apk_labels_new.append(apk_labels[apk_idx])
        apk_path_list_new.append(apk_path_list[apk_idx])
        count += 1
        print(apk_path_list[apk_idx] + " Done! " + str(count))
    except:
        print(apk_path_list[apk_idx] + " failed!")
        continue
    # if count == 100:
    #     break
    
time_point_1 = time.time()
    
apk_test_or_train_new = np.array(apk_test_or_train_new)
apk_labels_new = np.array(apk_labels_new)

sio.savemat("./output_dir/apk_list_and_label.mat", mdict={ "apk_labels":apk_labels_new, 
                                            "apk_test_or_train": apk_test_or_train_new,
                                            "apk_path_list": apk_path_list_new}, long_field_names=True)

try:
    with open("./output_dir/apk_list.pkl", 'wb') as f:
        pickle.dump(apk_list, f)
    print("apk_list save done!")
except Exception as e:
    print("apk_list save failed")
    print(e)

# judge 0 or 1 
def fun(feature, self_feature_list):
    if feature in self_feature_list:
        return 1
    else:
        return 0 

# Construct vector for every APK
apk_feature_list_per = []
apk_feature_list_cla = []
apk_feature_list_lib = []
apk_feature_list_int = []
apk_feature_list_met = []
count = 0
time_point_2 = time.time()
for apk in apk_list:
    #apk_feature = []
    ''' This part optimized
    apk.per_vector = np.array([fun(x, apk.permission) for x in list(per_dict.values())])
    apk.cla_vector = np.array([fun(x, apk.apk_class) for x in list(cla_dict.values())])
    apk.lib_vector = np.array([fun(x, apk.library) for x in list(lib_dict.values())])
    apk.int_vector = np.array([fun(x, apk.interface) for x in list(int_dict.values())])
    apk.met_vector = np.array([fun(x, apk.methods) for x in list(met_dict.values())])
    '''
    
    apk.per_vector = [0] * len(per_dict.values())
    for item in apk.permission:
        apk.per_vector[item] = 1
    
    apk.cla_vector = [0] * len(cla_dict.values())
    for item in apk.apk_class:
        apk.cla_vector[item] = 1
        
    apk.lib_vector = [0] * len(lib_dict.values())
    for item in apk.library:
        apk.lib_vector[item] = 1
        
    apk.int_vector = [0] * len(int_dict.values())
    for item in apk.interface:
        apk.int_vector[item] = 1
        
    apk.met_vector = [0] * len(met_dict.values())
    for item in apk.methods:
        apk.met_vector[item] = 1

    apk.per_vector = np.array(apk.per_vector, dtype=bool)
    apk.cla_vector = np.array(apk.cla_vector, dtype=bool)
    apk.lib_vector = np.array(apk.lib_vector, dtype=bool)
    apk.int_vector = np.array(apk.int_vector, dtype=bool)
    apk.met_vector = np.array(apk.met_vector, dtype=bool)
        
    # apk_feature.append(apk.per_vector)
    # apk_feature.append(apk.cla_vector)
    # apk_feature.append(apk.lib_vector)
    # apk_feature.append(apk.int_vector)
    # apk_feature.append(apk.met_vector)
    
    apk_feature_list_per.append(apk.per_vector)
    apk_feature_list_cla.append(apk.cla_vector)
    apk_feature_list_lib.append(apk.lib_vector)
    apk_feature_list_int.append(apk.int_vector)
    apk_feature_list_met.append(apk.met_vector)
    
    count += 1
    print(str(count) + "apk vector build")

apk_feature_list_per = np.array(apk_feature_list_per, dtype=bool)
apk_feature_list_cla = np.array(apk_feature_list_cla, dtype=bool)
apk_feature_list_lib = np.array(apk_feature_list_lib, dtype=bool)
apk_feature_list_int = np.array(apk_feature_list_int, dtype=bool)
apk_feature_list_met = np.array(apk_feature_list_met, dtype=bool)

apk_feature_list_per = sparse.csr_matrix(apk_feature_list_per)
apk_feature_list_cla = sparse.csr_matrix(apk_feature_list_cla)
apk_feature_list_lib = sparse.csr_matrix(apk_feature_list_lib)
apk_feature_list_int = sparse.csr_matrix(apk_feature_list_int)
apk_feature_list_met = sparse.csr_matrix(apk_feature_list_met)

time_point_3 = time.time()


try:
    sio.savemat("./output_dir/apk_feature_list_per.mat", 
                mdict={"apk_feature_list_per": apk_feature_list_per})
    print("apk_feature_list_per.mat save done!")
except Exception as e:
    print("apk_feature_list_per.mat save faild")
    print(e)
    
try:
    sio.savemat("./output_dir/apk_feature_list_cla.mat", 
                mdict={"apk_feature_list_per": apk_feature_list_cla})
    print("apk_feature_list_cla.mat save done!")
except Exception as e:
    print("apk_feature_list_cla.mat save faild")
    print(e)
    
try:
    sio.savemat("./output_dir/apk_feature_list_lib.mat", 
                mdict={"apk_feature_list_lib": apk_feature_list_lib})
    print("apk_feature_list_lib.mat save done!")
except Exception as e:
    print("apk_feature_list_lib.mat save faild")
    print(e)
    
try:
    sio.savemat("./output_dir/apk_feature_list_int.mat", 
                mdict={"apk_feature_list_int": apk_feature_list_int})
    print("apk_feature_list_int.mat save done!")
except Exception as e:
    print("apk_feature_list_int.mat save faild")
    print(e)
    
try:
    sio.savemat("./output_dir/apk_feature_list_met.mat", 
                mdict={"apk_feature_list_met": apk_feature_list_met})
    print("apk_feature_list_met.mat save done!")
except Exception as e:
    print("apk_feature_list_met.mat save faild")
    print(e)
    
    
    # np.save("./output_dir/per_dict.npy", per_dict)
    # np.save("./output_dir/cla_dict.npy", cla_dict)
    # np.save("./output_dir/lib_dict.npy", lib_dict)
    # np.save("./output_dir/int_dict.npy", int_dict)
    # np.save("./output_dir/met_dict.npy", met_dict)

try:
    with open("./output_dir/per_dict.txt", 'w') as f:
        for i in per_dict.items():
            f.write(str(i)+'\n')
    print("per_dict save done")
except Exception as e:
    print("per_dict save failed")
    print(e)
    
try:
    with open("./output_dir/cla_dict.txt", 'w') as f:
        for i in cla_dict.items():
            f.write(str(i)+'\n')
    print("cla_dict save done")
except Exception as e:
    print("cla_dict save failed")
    print(e)
    
try:
    with open("./output_dir/lib_dict.txt", 'w') as f:
        for i in lib_dict.items():
            f.write(str(i)+'\n')
    print("lib_dict save done")
except Exception as e:
    print("lib_dict save failed")
    print(e)
    
try:
    with open("./output_dir/int_dict.txt", 'w') as f:
        for i in int_dict.items():
            f.write(str(i)+'\n')
    print("int_dict save done")
except Exception as e:
    print("int_dict save failed")
    print(e)
    
try:
    with open("./output_dir/met_dict.txt", 'w') as f:
        for i in met_dict.items():
            f.write(str(i)+'\n')
    print("met_dict save done")
except Exception as e:
    print("met_dict save failed")
    print(e)
        


# Construct graph
permission_graph = np.zeros((len(apk_list), len(apk_list)), dtype=bool)
class_graph = np.zeros((len(apk_list), len(apk_list)), dtype=bool)
library_graph = np.zeros((len(apk_list), len(apk_list)), dtype=bool)
interface_graph = np.zeros((len(apk_list), len(apk_list)), dtype=bool)
method_graph = np.zeros((len(apk_list), len(apk_list)), dtype=bool)

count = 0
time_point_4 = time.time()
for i in range(len(apk_list)):
    for j in range(i, len(apk_list)):
        multi_result = np.multiply(apk_list[i].per_vector, apk_list[j].per_vector)
        if (1 in multi_result) or i==j:
            permission_graph[i][j] = 1
            
        multi_result = np.multiply(apk_list[i].cla_vector, apk_list[j].cla_vector)
        if (1 in multi_result) or i==j:
            class_graph[i][j] = 1
            
        multi_result = np.multiply(apk_list[i].lib_vector, apk_list[j].lib_vector)
        if (1 in multi_result) or i==j:
            library_graph[i][j] = 1
            
        multi_result = np.multiply(apk_list[i].int_vector, apk_list[j].int_vector)
        if (1 in multi_result) or i==j:
            interface_graph[i][j] = 1
            
        multi_result = np.multiply(apk_list[i].met_vector, apk_list[j].met_vector)
        if (1 in multi_result) or i==j:
            method_graph[i][j] = 1
    count += 1
    print(str(count) + "apk graph build")

# Get the symmetric matrix. Bool array, no need to subtract np.eye(length)
permission_graph = permission_graph + permission_graph.T
class_graph = class_graph + class_graph.T
library_graph = library_graph + library_graph.T
interface_graph = interface_graph + interface_graph.T
method_graph = method_graph + method_graph.T


permission_graph = sparse.csr_matrix(permission_graph | np.eye(len(apk_list), dtype=bool))
class_graph = sparse.csr_matrix(class_graph | np.eye(len(apk_list), dtype=bool))
library_graph = sparse.csr_matrix(library_graph | np.eye(len(apk_list), dtype=bool))
interface_graph = sparse.csr_matrix(interface_graph | np.eye(len(apk_list), dtype=bool))
method_graph = sparse.csr_matrix(method_graph | np.eye(len(apk_list), dtype=bool))

time_point_5 = time.time()

sio.savemat("./output_dir/apk_graph.mat",mdict={"permission_graph":permission_graph,
                                 "class_graph":class_graph,
                                 "library_graph":library_graph,
                                 "interface_graph":interface_graph,
                                 "method_graph":method_graph}, long_field_names=True)
print("Graph save done!")

print(time_point_0, time_point_1, time_point_2, time_point_3, time_point_4, time_point_5)

    
